import subprocess
import os
import tempfile
import logging
import shutil
import json
import librosa
import numpy as np
import sys

# 设置环境变量阻止使用MXNet
os.environ["TRANSFORMERS_NO_MXNET"] = "1"
os.environ["KERAS_BACKEND"] = "tensorflow"  # 强制Keras使用TensorFlow后端
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # 减少TensorFlow警告

# 检查ffmpeg并设置环境变量
def setup_ffmpeg():
    """检查并设置ffmpeg路径"""
    ffmpeg_found = False
    
    # 首先检查项目根目录下的ffmpeg_bin文件夹
    ffmpeg_bin_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg_bin", "ffmpeg.exe")
    if os.path.exists(ffmpeg_bin_path):
        os.environ["PATH"] = ffmpeg_bin_path.replace("ffmpeg.exe", "") + os.pathsep + os.environ["PATH"]
        os.environ["FFMPEG_BINARY"] = ffmpeg_bin_path
        print(f"使用FFmpeg: {ffmpeg_bin_path}")
        ffmpeg_found = True
    else:
        # 检查项目根目录
        ffmpeg_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg.exe")
        if os.path.exists(ffmpeg_path):
            os.environ["PATH"] = os.path.dirname(os.path.abspath(__file__)) + os.pathsep + os.environ["PATH"]
            os.environ["FFMPEG_BINARY"] = ffmpeg_path
            print(f"使用FFmpeg: {ffmpeg_path}")
            ffmpeg_found = True
    
    # 检查系统PATH中是否有ffmpeg
    if not ffmpeg_found:
        try:
            subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            print("系统PATH中找到ffmpeg")
            ffmpeg_found = True
        except Exception:
            print("警告: 未找到ffmpeg，音频处理功能可能受限")
    
    return ffmpeg_found

# 设置ffmpeg路径
setup_ffmpeg()

# 确保在导入其他库之前设置环境变量
try:
    from transformers import pipeline
    import opencc
    import ffmpeg
except ImportError as e:
    print(f"导入错误: {e}")
except Exception as e:
    print(f"加载模型时发生错误: {e}")

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AudioExtractorTranscriber:
    def __init__(self, model_name="openai/whisper-small"):
        """初始化音频提取和转录器"""
        # 确保ffmpeg设置正确
        setup_ffmpeg()
        
        try:
            self.transcriber = pipeline("automatic-speech-recognition", 
                                   model=model_name,
                                   return_timestamps=True)
            logger.info("成功初始化Whisper转录模型")
        except Exception as e:
            logger.error(f"初始化Whisper模型失败: {str(e)}")
            self.transcriber = None
            
        # 初始化繁体到简体转换器
        try:
            self.converter = opencc.OpenCC('t2s')  # 't2s'表示繁体到简体
            logger.info("成功初始化繁体到简体转换器")
        except Exception as e:
            logger.error(f"初始化繁体到简体转换器失败: {str(e)}")
            self.converter = None
    
    def transcribe_video(self, video_path, output_path=None):
        """
        转录视频文件，并可选择性地保存转录文本到文件

        参数:
            video_path (str): 视频文件路径
            output_path (str, optional): 转录结果保存路径，如不提供则不保存文件
            
        返回:
            str: 转录文本，如果处理失败则返回空字符串
        """
        # 提取音频
        audio_path = self.extract_audio(video_path)
        
        # 如果音频提取失败，尝试使用替代方法
        if not audio_path:
            logger.warning(f"使用ffmpeg-python提取音频失败，尝试使用subprocess方法")
            audio_path = self.extract_audio_subprocess(video_path)
            
            if not audio_path:
                logger.error(f"所有音频提取方法都失败，无法处理视频: {video_path}")
                return ""
        
        # 音频转录
        try:
            # 如果转录器不可用，返回空结果
            if self.transcriber is None:
                logger.warning("转录器不可用，将返回空结果")
                return ""
                
            logger.info(f"开始转录音频: {audio_path}")
            result = self.transcriber(audio_path)
            logger.info(f"转录结果: {result}")
            
            # 组织文本，并转换为简体中文
            text = result.get("text", "")
            if self.converter:
                text = self.converter.convert(text)
            
            # 保存转录结果
            if output_path and text:
                output_dir = os.path.dirname(output_path)
                if output_dir and not os.path.exists(output_dir):
                    os.makedirs(output_dir, exist_ok=True)
                
                try:
                    with open(output_path, 'w', encoding='utf-8') as f:
                        f.write(text)
                    logger.info(f"转录文本已保存到: {output_path}")
                except Exception as e:
                    logger.error(f"保存转录文本失败: {str(e)}")
            
            return text
            
        except Exception as e:
            logger.error(f"转录音频时出错: {str(e)}")
            import traceback
            logger.error(traceback.format_exc())
            return ""
        # 我们不再删除音频文件，而是保留它们
    
    def process(self, input_data):
        """
        处理单个视频，提取音频并转录。
        
        参数:
            input_data (dict): 包含视频路径的字典
            
        返回:
            dict: 包含转录结果的字典
        """
        video_path = input_data['video_path']
        
        # 提取音频
        audio_path = self.extract_audio(video_path)
        
        # 如果音频提取失败，尝试使用替代方法
        if not audio_path:
            logger.warning(f"使用ffmpeg-python提取音频失败，尝试使用subprocess方法")
            audio_path = self.extract_audio_subprocess(video_path)
            
            if not audio_path:
                logger.error(f"所有音频提取方法都失败，无法处理视频: {video_path}")
                return {
                    'audio_path': None,
                    'transcript': "",
                    'chunks': []
                }
        
        try:
            # 如果转录器不可用，返回空结果
            if self.transcriber is None:
                logger.warning("转录器不可用，将返回空结果")
                return {
                    'audio_path': audio_path,
                    'transcript': "",
                    'chunks': []
                }
                
            logger.info(f"开始转录音频: {audio_path}")
            result = self.transcriber(audio_path)
            logger.info(f"转录结果: {result}")
            
            # 组织时间戳和文本，并转换为简体中文
            text = result.get("text", "")
            if self.converter:
                text = self.converter.convert(text)
            
            chunks = []
            if "chunks" in result:
                for chunk in result["chunks"]:
                    chunk_text = chunk.get("text", "")
                    if self.converter:
                        chunk_text = self.converter.convert(chunk_text)
                    chunks.append({
                        "text": chunk_text,
                        "start": chunk["timestamp"][0] if "timestamp" in chunk else 0,
                        "end": chunk["timestamp"][1] if "timestamp" in chunk else 1
                    })
            
            # 如果没有chunks，创建一个默认chunk
            if not chunks and text:
                chunks.append({
                    "text": text,
                    "start": 0,
                    "end": 1
                })
            
            return {
                'audio_path': audio_path,
                'transcript': text,
                'chunks': chunks
            }
        except Exception as e:
            logger.error(f"转录音频时出错: {str(e)}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                'audio_path': audio_path,
                'transcript': "",
                'chunks': []
            }
    
    def extract_audio(self, video_path):
        """使用ffmpeg-python从视频中提取音频"""
        try:
            # 检查项目根目录下的ffmpeg_bin文件夹是否有ffmpeg
            ffmpeg_bin_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg_bin", "ffmpeg.exe")
            if os.path.exists(ffmpeg_bin_path):
                os.environ["FFMPEG_BINARY"] = ffmpeg_bin_path
                logger.info(f"使用FFmpeg: {ffmpeg_bin_path}")
            else:
                # 向后兼容，检查老路径
                ffmpeg_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg.exe")
                if os.path.exists(ffmpeg_path):
                    os.environ["FFMPEG_BINARY"] = ffmpeg_path
                    logger.info(f"使用FFmpeg: {ffmpeg_path}")
            
            # 创建音频文件目录
            audio_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "audio_files")
            os.makedirs(audio_dir, exist_ok=True)
            
            # 使用项目目录中的audio_files文件夹存放音频
            audio_path = os.path.join(audio_dir, f"{os.path.basename(video_path)}.wav")
            
            try:
                # 尝试使用ffmpeg-python
                (
                    ffmpeg
                    .input(video_path)
                    .output(audio_path, acodec='pcm_s16le', ac=1, ar='16k')
                    .overwrite_output()
                    .run(capture_stdout=True, capture_stderr=True)
                )
                
                if os.path.exists(audio_path) and os.path.getsize(audio_path) > 0:
                    logger.info(f"成功提取音频到: {audio_path}")
                    return audio_path
                else:
                    logger.warning(f"音频文件不存在或为空: {audio_path}")
                    return None
            except ffmpeg.Error as e:
                stderr = e.stderr.decode() if hasattr(e, 'stderr') and e.stderr else "未知错误"
                logger.error(f"使用ffmpeg-python提取音频失败: {stderr}")
                return None
                
        except Exception as e:
            logger.error(f"使用ffmpeg-python提取音频失败: {str(e)}")
            return None
    
    def extract_audio_subprocess(self, video_path):
        """使用subprocess调用ffmpeg从视频中提取音频（备用方法）"""
        try:
            # 创建音频文件目录
            audio_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "audio_files")
            os.makedirs(audio_dir, exist_ok=True)
            
            # 使用项目目录中的audio_files文件夹存放音频
            audio_path = os.path.join(audio_dir, f"backup_{os.path.basename(video_path)}.wav")
            
            # 首先尝试使用已知的ffmpeg_bin中的ffmpeg.exe
            ffmpeg_bin_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg_bin", "ffmpeg.exe")
            
            # 如果ffmpeg_bin目录存在ffmpeg.exe，则使用它
            if os.path.exists(ffmpeg_bin_path):
                ffmpeg_path = ffmpeg_bin_path
            else:
                # 然后尝试项目根目录中的ffmpeg
                ffmpeg_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg.exe")
                
                # 如果项目目录中没有ffmpeg，则尝试使用系统PATH中的ffmpeg
                if not os.path.exists(ffmpeg_path):
                    ffmpeg_path = "ffmpeg"
            
            logger.info(f"使用ffmpeg路径: {ffmpeg_path}")
            
            command = [
                ffmpeg_path,
                "-i", video_path,
                "-q:a", "0",
                "-map", "a",
                audio_path,
                "-y"
            ]
            
            # 使用subprocess.PIPE捕获输出，并指定stderr=subprocess.PIPE避免解码错误
            process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                # 显式指定错误处理策略
                errors='ignore',
                # 使用通用UTF-8编码而非GBK
                encoding='utf-8'
            )
            
            # 等待进程完成，但不读取输出流（避免解码错误）
            process.wait()
            
            if process.returncode != 0:
                logger.error(f"ffmpeg命令执行失败，返回代码: {process.returncode}")
                return None
            
            if os.path.exists(audio_path) and os.path.getsize(audio_path) > 0:
                logger.info(f"使用subprocess成功提取音频到: {audio_path}")
                return audio_path
            else:
                logger.error("音频提取失败，文件不存在或为空")
                return None
                
        except Exception as e:
            logger.error(f"使用subprocess提取音频时出错: {str(e)}")
            return None
    
    def process_folder(self, folder_path, output_folder="transcripts"):
        """
        处理文件夹中的所有视频文件。
        
        参数:
            folder_path (str): 包含视频文件的文件夹路径
            output_folder (str): 输出转录结果的文件夹路径
            
        返回:
            dict: 视频ID到转录结果的映射
        """
        import json
        
        # 确保输出文件夹存在
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        
        results = {}
        
        # 获取所有视频文件
        video_files = [f for f in os.listdir(folder_path) 
                      if f.lower().endswith(('.mp4', '.avi', '.mov', '.mkv'))]
        
        for video_file in video_files:
            video_path = os.path.join(folder_path, video_file)
            video_id = os.path.splitext(video_file)[0]
            
            # 处理单个视频
            result = self.process({'video_path': video_path})
            
            # 保存转录结果
            output_path = os.path.join(output_folder, f"{video_id}.json")
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            
            results[video_id] = result
        
        return results

# 如果直接运行这个文件
if __name__ == "__main__":
    # 创建转录器实例
    transcriber = AudioExtractorTranscriber()
    
    # 设置视频文件夹路径和输出文件夹
    videos_folder = "videos"
    transcripts_folder = "transcripts"
    
    # 处理文件夹中的所有视频
    print(f"开始处理文件夹: {videos_folder}")
    results = transcriber.process_folder(videos_folder, transcripts_folder)
    
    # 输出结果信息
    print("\n处理结果汇总:")
    for video_name, result in results.items():
        print(f"视频: {video_name}")
        print(f"  转录文件: {result.get('transcript', '')}")
        print(f"  时间戳文件: ")
        for chunk in result.get('chunks', []):
            print(f"    [{chunk['start']:.2f} --> {chunk['end']:.2f}] {chunk['text']}")
        print()